# -*- coding: utf-8 -*-

import json
import random
import pandas as pd

import factory
import utils
import wqb

class Generator:
    """
    Alpha生成器
    """
    def __init__(self,wqbs:wqb.WQBSession):
        """
        初始化
        """
        self.wqbs = wqbs

    def process_datafields(self, df:pd.DataFrame):
        """处理数据字段"""
        datafields = []
        datafields += df[df['type'] == "MATRIX"]["id"].tolist()
        datafields += self.get_vec_fields(df[df['type'] == "VECTOR"]["id"].tolist())
        # return ["winsorize(ts_backfill(%s, 120), std=4)" % field for field in datafields]
        return datafields
    
    def get_vec_fields(self,fields):
        vec_ops = ["vec_avg", "vec_sum"]
        vec_fields = []
    
        for field in fields:
            for vec_op in vec_ops:
                if vec_op == "vec_choose":
                    vec_fields.append("%s(%s, nth=-1)"%(vec_op, field))
                    vec_fields.append("%s(%s, nth=0)"%(vec_op, field))
                else:
                    vec_fields.append("%s(%s)"%(vec_op, field))
    
        return(vec_fields)
    
    def generate_first_with_fields(self, fields:list, settings:dict) -> list:
        print(f"📋 根据指定字段列表构建...")
        first_order = factory.first_order_factory(fields, factory.ts_ops)
        print(f'📋 构建结束，共{len(first_order)}个表达式, 前五个表达式如下: \n{first_order[:5]}')
        sim_data_list =  factory.generate_sim_data(first_order, settings)
        print(f'📋 生成结束，共{len(sim_data_list)}个alpha...')
        print(f'📋 开始保存alpha...')
        # self.mapper.bath_save(sim_data_list)
        print(f'📋 保存结束...')
        return sim_data_list
    
    def generate_first(self, dataset_id:str, settings:dict) -> list:
        print(f"📋 获取数据集{dataset_id}字段列表...")
        # 1. 获取数据集字段
        fields = utils.get_dataset_fields(
            self.wqbs
            , dataset_id=dataset_id,
            region=settings['region'],
            universe=settings['universe'],
            delay = settings['delay']
        )

        print(f'📋 数据集{dataset_id}共{len(fields)}个字段...')
        if  len(fields) == 0:
            print(f"📋 数据集{dataset_id}字段列表为空，请检查数据集是否正确...")
            return []
        fields = fields[:200]
        # 2. 处理数据字段
        df = pd.DataFrame(fields)
        print(f'📋 开始处理字段...')
        pc_fields = self.process_datafields(df)
        print(f'📋 处理结束，共{len(pc_fields)}个字段...')
        print(f'📋 开始构建表达式...')
        # 3. 构建表达式
        # first_order = factory.first_order_factory(pc_fields, factory.ts_ops)
        first_order = factory.first_order_factory_new(pc_fields)
        #prefix= dataset_config.get_dataset_config(dataset_id)['field_prefix']
        sim_data_list =  factory.generate_sim_data(first_order, settings)
        print(f'📋 生成结束，共{len(sim_data_list)}个alpha...')
        utils.save_to_file(sim_data_list, region=settings['region'], dataset_id=dataset_id)